Skip to main content
Version: Next

Multiclass SHAP values

A multiclass LightGBM prediction has one output score per class. Its featuresShapCol contains contributions for every feature and every class, plus a bias value for each class, packed into one Spark vector. It is not one importance number per class.

This tutorial uses objective="multiclass" with five classes and a small, entirely synthetic dataset containing numeric and one-hot encoded categorical features. It explains the vector layout and calculates per-class importance over all explained rows. These are explanations of model scores, not accuracy measurements or causal effects.

Run this notebook in a Spark session with SynapseML installed. It reuses the session and does not change its configuration or stop it. See the LightGBM overview for installation and model parameters.

For objective="binary", the native booster reports one class and returns only the positive-class contributions, with F + 1 entries. This notebook does not cover a two-label model trained with objective="multiclass", or the multiclassova objective.

import numpy as np

from pyspark.ml import Pipeline
from pyspark.ml.feature import OneHotEncoder, StringIndexer, VectorAssembler
from pyspark.ml.functions import vector_to_array
from pyspark.sql import SparkSession, functions as sf

from synapse.ml.lightgbm import LightGBMClassifier

spark = SparkSession.builder.getOrCreate()
n_classes = 5
if n_classes <= 2:
    raise ValueError("This example requires more than two multiclass labels.")

# A bounded teaching dataset, not a sample from a real population.
n_rows = 201
rows = [
    (
        i,
        float(i % 5) + 0.1 * (i % 3),
        float((7 * i) % 19),
        f"segment-{i % 4}",
        float(((i % 5) + (1 if i % 4 >= 2 else 0)) % n_classes),
    )
    for i in range(n_rows)
]
data = spark.createDataFrame(
    rows, "id long, numeric_a double, numeric_b double, segment string, label double"
).coalesce(1)
assert data.select("label").distinct().count() == n_classes

Preserve the encoded feature names​

The original data has three feature columns, but one-hot encoding expands segment into several vector slots. The number of original columns is therefore not the model's feature count.

Use the assembled vector metadata to name its slots. The assembler below places the two numeric slots first, followed by the categorical vector. Keeping that input order explicit also lets us group the categorical slots later.

numeric_columns = ["numeric_a", "numeric_b"]
indexer = StringIndexer(
    inputCol="segment", outputCol="segment_index", stringOrderType="alphabetAsc"
)
encoder = OneHotEncoder(
    inputCols=["segment_index"], outputCols=["segment_encoded"], dropLast=False
)
assembler = VectorAssembler(
    inputCols=numeric_columns + ["segment_encoded"], outputCol="features"
)
preprocessor = Pipeline(stages=[indexer, encoder, assembler]).fit(data)
prepared = preprocessor.transform(data)

model = LightGBMClassifier(
    objective="multiclass",
    featuresCol="features",
    labelCol="label",
    featuresShapCol="featuresShap",
    numIterations=10,
    numLeaves=8,
    minDataInLeaf=1,
    numTasks=1,
    numThreads=1,
    seed=29,
).fit(prepared)

n_classes = model.getBoosterNumClasses()
n_features = model.getBoosterNumFeatures()
metadata = prepared.schema["features"].metadata["ml_attr"]
attributes = sorted(
    [attribute for group in metadata["attrs"].values() for attribute in group],
    key=lambda attribute: attribute["idx"],
)
if n_features != metadata["num_attrs"] or [
    attribute["idx"] for attribute in attributes
] != list(range(n_features)):
    raise ValueError(
        "Model and assembled feature dimensions differ; do not guess slot names."
    )
slot_names = [attribute["name"] for attribute in attributes]
assert slot_names[:2] == numeric_columns
assert n_classes == 5 and n_features + 1 != n_classes
assert n_features == len(numeric_columns) + preprocessor.stages[1].categorySizes[0]
print(list(enumerate(slot_names)))

Read the class blocks and the bias​

Let K = model.getBoosterNumClasses() and F = model.getBoosterNumFeatures(). For this multiclass objective with K > 2, each row contains K * (F + 1) values in class-major order:

[class 0: feature 0, ..., feature F-1, bias,
 class 1: feature 0, ..., feature F-1, bias,
 ...]

Class block k corresponds to the model's numeric label k; our labels are explicitly 0, ..., K-1. If your target was indexed, retain that fitted label indexer's mapping to recover the original names.

The last slot of each block is the expected value (bias), constant across rows for the fitted model. The feature contributions plus that bias reconstruct the corresponding rawPrediction score. They do not add up to a probability. Exclude bias when ranking features, but keep it when checking reconstruction.

The checks below inspect at most 20 rows and use an absolute reconstruction tolerance of 1e-6 * max(1, abs(raw score)) to allow floating-point accumulation differences. Different platforms can produce different fitted scores; no particular score or ranking is expected.

predictions = model.transform(prepared)
flat_predictions = predictions.withColumn("shap_array", vector_to_array("featuresShap"))
expected_size = n_classes * (n_features + 1)
bad_values = sf.exists(
    sf.col("shap_array"),
    lambda value: value.isNull()
    | sf.isnan(value)
    | (sf.abs(value) == sf.lit(float("inf"))),
)
assert (
    flat_predictions.where(
        sf.col("shap_array").isNull()
        | (sf.size("shap_array") != expected_size)
        | bad_values
    ).count()
    == 0
)

inspected = (
    predictions.select("id", "rawPrediction", "featuresShap")
    .orderBy("id")
    .limit(20)
    .collect()
)
flat_values = np.asarray([row.featuresShap.toArray() for row in inspected])
raw_scores = np.asarray([row.rawPrediction.toArray() for row in inspected])
class_blocks = flat_values.reshape(len(inspected), n_classes, n_features + 1)

reconstructed = class_blocks.sum(axis=2)
tolerance = 1e-6 * np.maximum(1.0, np.abs(raw_scores))
assert np.all(np.abs(reconstructed - raw_scores) <= tolerance)
bias = class_blocks[:, :, -1]
assert np.all(np.abs(bias - bias[:1]) <= 1e-6 * np.maximum(1.0, np.abs(bias[:1])))

# This deliberately wrong ordering must not pass on the asymmetric example.
wrong_blocks = flat_values.reshape(len(inspected), n_features + 1, n_classes).transpose(
    0, 2, 1
)
assert not np.all(np.abs(wrong_blocks.sum(axis=2) - raw_scores) <= tolerance)
print("Class-block shape:", class_blocks.shape)
print("Per-row raw-score reconstruction and constant bias checks passed.")

Mean absolute contribution for each output class​

For output class k and feature slot j, calculate the mean of abs(SHAP[k, j]) over all rows being explained. This answers which slots influence that class's output on this dataset. Restricting the rows to those whose true label or prediction is k would answer a different question.

A signed mean can cancel positive and negative contributions. Mean absolute SHAP is also different from the model-wide split/gain importance returned by getFeatureImportances().

The aggregation stays in Spark; only the small aggregate row is collected. The bias slot is excluded.

slot_expressions = [
    sf.avg(sf.abs(sf.col("shap_array")[class_id * (n_features + 1) + slot])).alias(
        f"class_{class_id}_slot_{slot}"
    )
    for class_id in range(n_classes)
    for slot in range(n_features)
]
slot_means = flat_predictions.agg(*slot_expressions).first()
slot_summary_rows = [
    (class_id, slot_names[slot], float(slot_means[f"class_{class_id}_slot_{slot}"]))
    for class_id in range(n_classes)
    for slot in range(n_features)
]
slot_summary = spark.createDataFrame(
    slot_summary_rows, "output_class int, encoded_feature string, mean_abs_shap double"
)
slot_summary.orderBy("output_class", sf.desc("mean_abs_shap")).show(
    n_classes * n_features, truncate=False
)

Summarize the original categorical column​

To combine the one-hot slots for segment, first sum their signed contributions within each row and output class, then take the absolute value and average over rows. Taking the absolute value before combining the slots gives a different statistic.

This is an additive summary of the existing encoded-feature explanation. It does not recompute SHAP with the original categorical column treated as a single coalition player. The numeric columns each keep their single slot. The mapping below follows this notebook's explicit assembler order and checked encoder width.

feature_groups = {
    "numeric_a": [0],
    "numeric_b": [1],
    "segment": list(range(len(numeric_columns), n_features)),
}
assert feature_groups["segment"]
group_expressions = []
for class_id in range(n_classes):
    for group_name, slots in feature_groups.items():
        signed_contribution = sum(
            (
                sf.col("shap_array")[class_id * (n_features + 1) + slot]
                for slot in slots
            ),
            sf.lit(0.0),
        )
        group_expressions.append(
            sf.avg(sf.abs(signed_contribution)).alias(f"class_{class_id}_{group_name}")
        )
group_means = flat_predictions.agg(*group_expressions).first()
group_summary_rows = [
    (class_id, group_name, float(group_means[f"class_{class_id}_{group_name}"]))
    for class_id in range(n_classes)
    for group_name in feature_groups
]
spark.createDataFrame(
    group_summary_rows,
    "output_class int, original_feature string, mean_abs_shap double",
).orderBy("output_class", sf.desc("mean_abs_shap")).show(
    n_classes * len(feature_groups), truncate=False
)

Check the aggregates on this bounded example​

For this 201-row synthetic dataset only, we can compare the Spark aggregates with a direct NumPy calculation. Do not collect an unbounded production dataset to do this: the Spark aggregation above already computes the summaries without collecting all rows.

These checks validate the class/feature indexing and the order of categorical aggregation. They do not establish that a particular feature causes an outcome or that a model is accurate.

reference_rows = predictions.select("featuresShap").limit(n_rows + 1).collect()
assert len(reference_rows) == n_rows
reference = np.asarray([row.featuresShap.toArray() for row in reference_rows]).reshape(
    n_rows, n_classes, n_features + 1
)
expected_slot_means = np.abs(reference[:, :, :-1]).mean(axis=0)
for class_id in range(n_classes):
    for slot in range(n_features):
        np.testing.assert_allclose(
            slot_means[f"class_{class_id}_slot_{slot}"],
            expected_slot_means[class_id, slot],
            rtol=1e-9,
            atol=1e-10,
        )
    for group_name, slots in feature_groups.items():
        expected = np.abs(reference[:, class_id, slots].sum(axis=1)).mean()
        np.testing.assert_allclose(
            group_means[f"class_{class_id}_{group_name}"],
            expected,
            rtol=1e-9,
            atol=1e-10,
        )
print("Encoded-slot and original-column aggregate checks passed.")